#!/bin/bash

set -e

source "$(
  cd "$(dirname "$0")" || exit 1
  pwd
)/../common.sh"

#下载依赖和工具包
function install_development() {
  #sudo yum install -y wget tar unzip patch 
  source /etc/profile
}

#获取cuda
function download_cuda() {
  echo "# cuda-11.4正在下载"
  wget -t 40 -c -P "$1"\
    https://developer.download.nvidia.com/compute/cuda/11.4.0/local_installers/cuda_11.4.0_470.42.01_linux_sbsa.run
  #cp -r  "$2"/src/cuda_11.4.0_470.42.01_linux_sbsa.run "$1"/
  if [[ $? -eq 0 ]]; then
    echo -e "\033[1;32;1m# cuda-11.4拷贝成功\033[0m"
  else
    echo -e "\033[31m# cuda-11.4拷贝失败,请重新执行脚本\033[0m"
    exit 1
  fi
}

function install_cuda() {
  echo "# cuda-11.4正在安装"
  mkdir "$1"/tmpcuda
  sudo sh "$1"/cuda_11.4.0_470.42.01_linux_sbsa.run --tmpdir="$1"/tmpcuda  --override
  mkdir "$2"/cuda-11.4
  echo -e "\033[1;32;1m# cuda-11.4已经安装完毕\033[0m"
}

function write_cuda_module() {
  cat >"$1/cuda-11.4/cuda_modulefiles" <<EOF
#%Module1.0
conflict cuda
variable modfile [file normalize [info script]]
proc getModulefileDir {} {
  variable modfile
  set modfile_path [file dirname \$modfile]
  return \$modfile_path
}
prepend-path PATH \/usr/local/cuda-11.4/bin
prepend-path LD_LIBRARY_PATH \/usr/local/cuda-11.4/lib64
EOF
}

function main() {
    install_development
    download_cuda "$1" "$4"
    install_cuda "$1" "$2"
    write_cuda_module "$2"
    rm -rf "$1"/tmpcuda
}

main "$@"